import torch
import torch.nn as nn

def cal_power(x: torch.Tensor):
    """Calculate the average power for each example.

    This function assumes that the signal is structured as:

    .. math::

        Batch x Channel x IQ x Time.


    Args:
        x (torch.Tensor): Input Tensor (BxCxIQxT). x(t)=[i(t), q(t)], T=128

    Returns:
        [torch.Tensor]: power = 1/T * sum(i(t)**2+q(t)**2)
    """
    if len(x.shape) != 4:
        raise ValueError(
            "The inputs to the energy function must have 4 dimensions (BxCxIQxT), "
            "input shape was {}".format(x.shape)
        )
    if x.shape[2] != 2:
        raise ValueError(
            "The inputs to the energy function must be 'complex valued' by having 2 "
            "elements in the IQ dimension (BxCxIQxT), input shape was {}".format(
                x.shape
            )
        )
    iq_dim = 2
    time_dim = 3

    i, q = x.chunk(chunks=2, dim=iq_dim)
    power = (i * i) + (q * q)  # power is magnitude squared so sqrt cancels

    # pylint: disable=no-member
    # The linter isn't able to find the "mean" function but its there!
    x = torch.mean(power, dim=time_dim)

    # This Tensor still has an unnecessary singleton dimensions in IQ
    x = x.squeeze(dim=iq_dim)

    return x